package customExecutor;

/**
 * @Package: com.h2cloud.ThreadLocalTest.customExecutor
 * @ClassName: TransmittableThreadPoolExecutor
 * @Author: h2cloud
 * @CreateTime: 2023/12/21 11:08 下午
 * @Description:
 */

import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.*;

@Slf4j
public class TransmittableThreadPoolExecutor extends ThreadPoolExecutor {

	public TransmittableThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
		super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
	}

	public TransmittableThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
		super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
	}

	public TransmittableThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
		super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, handler);
	}

	public TransmittableThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
		super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
	}

	@Override
	public void execute(Runnable task) {
		super.execute(this.wrapRunnable(task));
	}

	@Override
	public <T> Future<T> submit(Runnable task, T result) {
		return super.submit(this.wrapRunnable(task), result);
	}

	@Override
	public <T> Future<T> submit(Callable<T> task) {
		return super.submit(this.wrapCallable(task));
	}

	@Override
	public Future<?> submit(Runnable task) {
		return super.submit(this.wrapRunnable(task));
	}

	// 包装 Runnable任务，将父线程的线程上下文变量复制到子线程中
	public Runnable wrapRunnable(final Runnable runnable) {
		final ITLCurrentUser.UserInfo userInfo = ITLCurrentUser.get();
		return () -> {
			ITLCurrentUser.UserInfo oldValue = ITLCurrentUser.get();
			ITLCurrentUser.set(userInfo);
			try {
				runnable.run();
			} catch (Exception e) {
				log.error("异常 ", e);
				throw e;
			} finally {
				ITLCurrentUser.set(oldValue);
			}
		};
	}

	public <T> Callable<T> wrapCallable(final Callable<T> callable) {
		final ITLCurrentUser.UserInfo userInfo = ITLCurrentUser.get();
		return () -> {
			ITLCurrentUser.set(userInfo);
			try {
				return callable.call();
			} catch (Exception e) {
				log.error("异常 ", e);
				throw e;
			} finally {
				ITLCurrentUser.clear();
			}
		};
	}
}
